#pragma once
#ifndef __EKF_H
#define __EKF_H

#include <opencv2/opencv.hpp>
#include <Eigen/Dense>
#include <ceres/ceres.h>
#include "tools.h"
#include "../../../tool/include/types.h"
// #include "../autoaim_node.hpp"

// 必要的工具函数
template <typename T>
void ypd_to_xyz(T &ypd, T &xyz)
{
    xyz[0] = ypd[2] * ceres::sin(M_PI / 2 - ypd[1]) * ceres::cos(M_PI / 2 + ypd[0]);
    xyz[1] = ypd[2] * ceres::sin(M_PI / 2 - ypd[1]) * ceres::sin(M_PI / 2 + ypd[0]);
    xyz[2] = ypd[2] * ceres::cos(M_PI / 2 - ypd[1]);
}

template <typename T>
void xyz_to_ypd(T &xyz, T &ypd)
{
    ypd[0] = ceres::atan2(xyz[1], xyz[0]) - M_PI / 2;
    ypd[1] = M_PI / 2 - ceres::atan2(ceres::sqrt(xyz[0] * xyz[0] + xyz[1] * xyz[1]), xyz[2]);
    ypd[2] = ceres::sqrt(xyz[0] * xyz[0] + xyz[1] * xyz[1] + xyz[2] * xyz[2]);
}

template <int V_X, int V_Z>
class autoaim_ekf
{
public:
    using Matrix_zzd = Eigen::Matrix<double, V_Z, V_Z>;
    using Matrix_xxd = Eigen::Matrix<double, V_X, V_X>;
    using Matrix_zxd = Eigen::Matrix<double, V_Z, V_X>;
    using Matrix_xzd = Eigen::Matrix<double, V_X, V_Z>;
    using Matrix_x1d = Eigen::Matrix<double, V_X, 1>;
    using Matrix_z1d = Eigen::Matrix<double, V_Z, 1>;

    Matrix_xxd A = Matrix_xxd::Identity();
    Matrix_xxd P = Matrix_xxd::Identity();
    Matrix_xxd Q = Matrix_xxd::Zero();
    Matrix_zzd R = Matrix_zzd::Zero();
    Matrix_zxd H = Matrix_zxd::Zero();
    Matrix_xzd K = Matrix_xzd::Zero();
    Matrix_z1d hx = Matrix_z1d::Zero();
    Matrix_x1d X = Matrix_x1d::Zero();

    int armor_num = 4; // 装甲板数量

    std::string param_path;
    cv::Mat img_ekf; // 测试图像

    void load_param_path(const std::string &param_path)
    {
        this->param_path = param_path;
    }

    // ekf的初始化
    void init(const Matrix_z1d &Z)
    {
        Eigen::Matrix<double, 3, 1> ypd, xyz;
        ypd[0] = Z[0];
        ypd[1] = Z[1];
        ypd[2] = Z[2];
        ypd_to_xyz(ypd, xyz);

        double x = xyz[0], y = xyz[1], r = 0.2;
        xyz[0] = x + r * sin(Z[3]);
        xyz[1] = y - r * cos(Z[3]);

        X = Matrix_x1d::Zero();
        X[0] = xyz[0];
        X[2] = xyz[1];
        X[4] = xyz[2];
        X[6] = Z[3];
        X[8] = r;

        A = Matrix_xxd::Identity();
        P = Matrix_xxd::Identity();
        Q = Matrix_xxd::Zero();
        R = Matrix_zzd::Zero();
        H = Matrix_zxd::Zero();
        K = Matrix_xzd::Zero();
        hx = Matrix_z1d::Zero();
    }

    // ekf的5个步骤
    void update_prior_estimate(const Matrix_z1d &Z, const double &dt)
    {
        get_A(dt);
        X = A * X;
        if (X[8] < 0.2)
        {
            X[8] = 0.2;
        }
        else if (X[8] > 0.3)
        {
            X[8] = 0.3;
        }
        // X[8] = 0.25;

        double min_dis = __DBL_MAX__;
        int min_dis_index = 0;
        for (int i = 0; i < armor_num; ++i)
        {
            double dis = sqrt(pow(sin(Z[3]) - sin(X[6] + i * 2 * M_PI / armor_num), 2) + pow(cos(Z[3]) - cos(X[6] + i * 2 * M_PI / armor_num), 2));
            if (dis < min_dis)
            {
                min_dis = dis;
                min_dis_index = i;
            }
        }
        X[6] = X[6] + min_dis_index * 2 * M_PI / armor_num;
        // std::cout << "ekf_prior_theta = " << X[6] << std::endl;

        while (1)
        {
            if (X[6] > M_PI)
            {
                X[6] = X[6] - 2 * M_PI;
            }
            else if (X[6] <= -M_PI)
            {
                X[6] = X[6] + 2 * M_PI;
            }
            else
            {
                break;
            }
        }
    }

    void update_prior_error_cov_mat(const double &dt)
    {
        get_A(dt);
        get_Q(dt);

        P = A * P * A.transpose() + Q;
    }

    void update_kalman_gain(void)
    {
        get_H_and_hx();
        get_R();

        K = P * H.transpose() * (H * P * H.transpose() + R).inverse();
    }

    void update_posterior_estimate(const Matrix_z1d &Z)
    {
        Matrix_z1d Z_minus_hx = Z - hx;
        while (1)
        {
            if (Z_minus_hx[3] > M_PI)
            {
                Z_minus_hx[3] = Z_minus_hx[3] - 2 * M_PI;
            }
            else if (Z_minus_hx[3] <= -M_PI)
            {
                Z_minus_hx[3] = Z_minus_hx[3] + 2 * M_PI;
            }
            else
            {
                break;
            }
        }
        X = X + K * Z_minus_hx;
        while (1)
        {
            if (X[6] > M_PI)
            {
                X[6] = X[6] - 2 * M_PI;
            }
            else if (X[6] <= -M_PI)
            {
                X[6] = X[6] + 2 * M_PI;
            }
            else
            {
                break;
            }
        }
        // std::cout << " ekf_theta = " << X[6] ;
    }

    void update_posterior_error_cov_mat(void)
    {
        P = (Matrix_xxd::Identity() - K * H) * P;
    }

    void get_A(const double &dt)
    {
        A(0, 1) = dt;
        A(2, 3) = dt;
        A(4, 5) = dt;
        A(6, 7) = dt;
    }

    void get_Q(const double &dt)
    {
        cv::Mat Cov_mu_tmp;
        cv::FileStorage f;
        f.open(param_path, cv::FileStorage::READ);
        f["Cov_mu"] >> Cov_mu_tmp;
        f.release();

        Eigen::Matrix<double, 5, 5> Cov_mu = Eigen::Matrix<double, 5, 5>::Zero();
        for (int i = 0; i < 5; ++i)
        {
            Cov_mu(i, i) = Cov_mu_tmp.at<double>(i, 0);
        }

        Eigen::Matrix<double, V_X, 5> G = Eigen::Matrix<double, V_X, 5>::Zero();
        G(0, 0) = 0.5 * dt * dt;
        G(1, 0) = dt;
        G(2, 1) = 0.5 * dt * dt;
        G(3, 1) = dt;
        G(4, 2) = 0.5 * dt * dt;
        G(5, 2) = dt;
        G(6, 3) = 0.5 * dt * dt;
        G(7, 3) = dt;
        G(8, 4) = dt;

        Q = G * Cov_mu * G.transpose();
    }

    void get_R(void)
    {
        cv::Mat R_tmp;
        cv::FileStorage f;
        f.open(param_path, cv::FileStorage::READ);
        f["R"] >> R_tmp;
        f.release();

        for (int i = 0; i < V_Z; ++i)
        {
            R(i, i) = R_tmp.at<double>(i, 0);
        }
    }

    void get_H_and_hx(void)
    {
        // 初始化
        ceres::Jet<double, V_X> x[V_X];
        ceres::Jet<double, V_X> z[V_Z];
        for (int i = 0; i < V_X; ++i)
        {
            x[i].a = X(i, 0);
            x[i].v[i] = 1;
        }

        ceres::Jet<double, V_X> xyz_armor[3];
        xyz_armor[0] = x[0] - x[8] * ceres::sin(x[6]);
        xyz_armor[1] = x[2] + x[8] * ceres::cos(x[6]);
        xyz_armor[2] = x[4];

        ceres::Jet<double, V_X> ypd_armor[3];

        xyz_to_ypd(xyz_armor, ypd_armor);

        for (int i = 0; i < 3; ++i)
        {
            z[i] = ypd_armor[i];
        }
        z[3] = x[6];

        // 获取H、hx
        for (int i = 0; i < V_Z; ++i)
        {
            hx(i, 0) = z[i].a;
            for (int j = 0; j < V_X; ++j)
            {
                H(i, j) = z[i].v[j];
            }
        }
    }

    void draw_and_show_result(const int &frame_width, const int &frame_height, const double &half_w, const double &half_h, const cv::Mat &cameraMatrix, const cv::Mat offset, const Pose &pose)
    {
        // cv::Mat points[armor_num][4];
        cv::Mat **points = new cv::Mat *[armor_num];
        for (int i = 0; i < armor_num; ++i)
        {
            points[i] = new cv::Mat[4];
        }

        double armor_pitch = (3 == armor_num) ? -M_PI / 12 : M_PI / 12;

        for (int i = 0; i < armor_num; ++i)
        {
            points[i][0] = (cv::Mat_<double>(3, 1) << half_w, 0, half_h);
            points[i][1] = (cv::Mat_<double>(3, 1) << -half_w, 0, half_h);
            points[i][2] = (cv::Mat_<double>(3, 1) << -half_w, 0, -half_h);
            points[i][3] = (cv::Mat_<double>(3, 1) << half_w, 0, -half_h); // 为装甲板角点赋初始值

            cv::Mat P_0 = (cv::Mat_<double>(3, 3) << 1, 0, 0, 0, cos(armor_pitch), -sin(armor_pitch), 0, sin(armor_pitch), cos(armor_pitch));
            cv::Mat P_1 = (cv::Mat_<double>(3, 3) << cos(i * 2 * M_PI / armor_num), -sin(i * 2 * M_PI / armor_num), 0, sin(i * 2 * M_PI / armor_num), cos(i * 2 * M_PI / armor_num), 0, 0, 0, 1);
            for (int j = 0; j < 4; ++j) // 先绕x轴转15度，再绕z轴转i*90度或120度等
            {
                points[i][j] = P_1 * P_0 * points[i][j];
            }
        }

        for (int i = 0; i < armor_num; ++i) // 再加上半径
        {
            for (int j = 0; j < 4; ++j)
            {
                points[i][j].at<double>(0, 0) = points[i][j].at<double>(0, 0) - X[8] * sin(i * 2 * M_PI / armor_num);
                points[i][j].at<double>(1, 0) = points[i][j].at<double>(1, 0) + X[8] * cos(i * 2 * M_PI / armor_num);
            }
        }
        // for (int j = 0; j < 4; ++j)//再加上半径
        // {
        //     points[0][j].at<double>(1, 0) += X[8];
        //     points[1][j].at<double>(0, 0) -= X[8];
        //     points[2][j].at<double>(1, 0) -= X[8];
        //     points[3][j].at<double>(0, 0) += X[8];
        // }

        for (int i = 0; i < armor_num; ++i) // 再将整个装甲板绕z轴转X[6]度
        {
            cv::Mat P_2 = (cv::Mat_<double>(3, 3) << cos(X[6]), -sin(X[6]), 0, sin(X[6]), cos(X[6]), 0, 0, 0, 1);
            for (int j = 0; j < 4; ++j)
            {
                points[i][j] = P_2 * points[i][j];
            }
        }
        for (int i = 0; i < armor_num; ++i) // 再将整个机器人平移到目的位置
        {
            for (int j = 0; j < 4; ++j)
            {
                points[i][j].at<double>(0, 0) += X[0];
                points[i][j].at<double>(1, 0) += X[2];
                points[i][j].at<double>(2, 0) += X[4];
            }
        }
        for (int i = 0; i < armor_num; ++i) // 将所有点变换到像素坐标系
        {
            for (int j = 0; j < 4; ++j)
            {
                land_to_ptz(points[i][j], points[i][j], pose);
                ptz_to_camera(points[i][j], points[i][j], offset);
                camera_to_pixel(points[i][j], points[i][j], cameraMatrix);
            }
        }

        img_ekf = cv::Mat(frame_height, frame_width, CV_8UC3, cv::Scalar(256, 256, 256));
        for (int i = 0; i < armor_num; ++i) // 画图
        {
            for (int j = 0; j < 4; ++j)
            {
                cv::circle(img_ekf, cv::Point(points[i][j].at<double>(0, 0), points[i][j].at<double>(1, 0)), 1, cv::Scalar(0, 0, 255), 1, cv::LINE_AA);
                cv::line(img_ekf, cv::Point(points[i][j % 4].at<double>(0, 0), points[i][j % 4].at<double>(1, 0)), cv::Point(points[i][(j + 1) % 4].at<double>(0, 0), points[i][(j + 1) % 4].at<double>(1, 0)), cv::Scalar(0, 0, 255), 1, cv::LINE_AA);
            }
        }

        cv::imshow("img_ekf", img_ekf);
        cv::waitKey(1);

        for (int i = 0; i < armor_num; ++i)
        {
            delete[] points[i];
        }
        delete[] points;
    }

    // Matrix_x
};

#endif
